#include "partition.h"

#include <assert.h>
#include <stdlib.h>
#include <stdio.h>

static void initDisjSet_base(
        DisjointSetElem *array,
        int16_t size);

static void unionDisjSet_base(
        DisjointSetElem *array,
        uint32_t capacity,
        int16_t id1,
        int16_t id2);

static PT_FIND_RESULT findDisjSet_base(
        DisjointSetElem *array,
        uint32_t capacity,
        int16_t id1,
        int16_t id2);

static void linkDisjSet_base(
        DisjointSetElem *array,
        uint32_t capacity,
        int16_t root1,
        int16_t root2);

static int16_t rootDisjSet_base(
        DisjointSetElem *array,
        uint32_t capacity,
        int16_t id);

static void initDisjSet_base(DisjointSetElem *array, int16_t capacity)
{
    int16_t i = 0;
    for(i=0; i<capacity; ++i) {
        array[i].rootId = -1;
    }
}

static void unionDisjSet_base(DisjointSetElem *array, uint32_t capacity, int16_t id1, int16_t id2)
{
    int16_t root1 = 0;
    int16_t root2 = 0;

    assert(id1 != id2);

    root1 = rootDisjSet_base(array, capacity, id1);
    root2 = rootDisjSet_base(array, capacity, id2);

    assert(root1 > -1);
    assert(root2 > -1);

    linkDisjSet_base(array, capacity, root1, root2);
}

static PT_FIND_RESULT findDisjSet_base(DisjointSetElem *array, uint32_t capacity, int16_t id1, int16_t id2)
{
    int16_t root1 = 0;
    int16_t root2 = 0;

    if(id1 == id2) {
        return PT_SAME_SET;
    }

    root1 = rootDisjSet_base(array, capacity, id1);
    root2 = rootDisjSet_base(array, capacity, id2);


    assert(root1 > -1);
    assert(root2 > -1);

    if(root1 == root2) {
        return PT_SAME_SET;
    } else {
        return PT_DIFF_SET;
    }
}

static void linkDisjSet_base(DisjointSetElem *array, uint32_t capacity, int16_t root1, int16_t root2)
{
    assert(root1 > -1 && (uint32_t)root1 < capacity);
    assert(root2 > -1 && (uint32_t)root2 < capacity);
    assert(array[root1].rootId < 0);
    assert(array[root2].rootId < 0);

    if(root1 == root2) {
        return;
    }
    /* if root2 is deeper, make root2 new root of root1
     * else, make root1 the new root2, and if they have same depth,
     * then incease depth of root1
     */
    if(array[root2].rootId < array[root1].rootId) {
        array[root1].rootId = root2;
    } else {
        if(array[root1].rootId == array[root2].rootId) {
            --array[root1].rootId;
        }
        array[root2].rootId = root1;
    }
    return;
}

static int16_t rootDisjSet_base(DisjointSetElem *array, uint32_t capacity, int16_t id)
{
    int16_t root = 0;
    int16_t nextRoot = 0;
    int16_t tmp = 0;
    uint32_t cnt = 0;
    uint8_t flag = 0;

    root = id;
    nextRoot = array[root].rootId;
    while(nextRoot >= 0 && root != nextRoot) {
        root = nextRoot;
        nextRoot = array[root].rootId;
        if(++cnt >= capacity) {
            flag = 1;
            break;
        }
    }
    if(flag == 1) {
        return -2;
    }

    assert(root >= 0);

    tmp = id;
    nextRoot = array[tmp].rootId;
    cnt = 0;
    while(nextRoot >= 0 && root != nextRoot) {
        array[tmp].rootId = root;
        tmp = nextRoot;
        nextRoot = array[tmp].rootId;
        if(++cnt >= capacity) {
            flag = 1;
            break;
        }
    }

    if(flag == 1) {
        return -2;
    }

    return root;
}


/* ********************************************************************************
 */

static void initSet_L1024(DisjointSet_L1024 *set)
{
    if(set == NULL) {
        return;
    }
    initDisjSet_base(set->set, LEN_1024);
    set->capacity = LEN_1024;
}

static void union_L1024(DisjointSet_L1024 *set, int16_t id1, int16_t id2)
{
    if(set == NULL) {
        return;
    }
    if(set->capacity > LEN_1024) {
        return;
    }
    if(id1 < 0 || (uint32_t)id1 >= set->capacity) {
        return;
    }
    if(id2 < 0 || (uint32_t)id2 >= set->capacity) {
        return;
    }

    unionDisjSet_base(set->set, set->capacity, id1, id2);
    return;
}

static PT_FIND_RESULT find_L1024(DisjointSet_L1024 *set, int16_t id1, int16_t id2)
{
    if(set == NULL) {
        return PT_ERROR;
    }
    if(set->capacity > LEN_1024) {
        return PT_ERROR;
    }
    if(id1 < 0 || (uint32_t)id1 >= set->capacity) {
        return PT_ERROR;
    }
    if(id2 < 0 || (uint32_t)id2 >= set->capacity) {
        return PT_ERROR;
    }

    return findDisjSet_base(set->set, set->capacity, id1, id2);
}

static int16_t root_L1024(DisjointSet_L1024 *set, int16_t id)
{
    if(set == NULL) {
        return -2;
    }
    if(set->capacity > LEN_1024) {
        return -2;
    }
    if(id < 0 || (uint32_t)id >= set->capacity) {
        return -2;
    }

    return rootDisjSet_base(set->set, set->capacity, id);
}

const DisjSetInterface_L1024 iDisjSet_L1024 = {
    initSet_L1024,
    union_L1024,
    find_L1024,
    root_L1024
};

/* ********************************************************************************
 */

void parition(
        const float * const mat,
        const uint32_t row,
        const uint32_t col,
        const uint32_t stride,
        const float thres,
        DisjointSet_L1024 *setBuf,
        Cluster_L128 *clusters)
{
    uint32_t i, j = 0;
    uint32_t tmpSize = 0;
    int16_t tmpRoot = 0;
    const float *p = NULL;

    if(mat == NULL) {
        return;
    }
    if(setBuf == NULL) {
        return;
    }
    if(clusters == NULL) {
        return;
    }

    /* initialize set buffer */
    iDisjSet_L1024.initSet(setBuf);

    if(row + col > setBuf->capacity) {
        return;
    }

    if(col > stride) {
        return;
    }

    if(thres < 0) {
        return;
    }

    p = mat;
    for(i = 0; i<row; ++i) {
        for(j=0; j<col; ++j) {
            if(p[i * stride + j] >= thres) {
                continue;
            }
            iDisjSet_L1024.unionSet(setBuf, i, row+j);
        }
    }

    for(i=0; i<(row + col); ++i) {
        tmpRoot = iDisjSet_L1024.rootSet(setBuf, i);

        printf("rt(%d) - %d \n", i, tmpRoot);
        assert(tmpRoot > -1 && tmpRoot < row+col);

        tmpSize = clusters[tmpRoot].size;
        clusters[tmpRoot].array[tmpSize] = i;
        ++clusters[tmpRoot].size;
    }
    return;
}
